import os
import shutil
from tqdm import tqdm
import argparse
from typing import List
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime

def get_every_class_num(txt_path_train,txt_path_eval,folder_path):
    # 需修改，根据自己的类别，注意一一对应
    class_categories=['mandatory','prohibitory','warning']
    class_num = len(class_categories)  # 样本类别数
    class_list = [i for i in range(class_num)]
    class_num_list = [0 for i in range(class_num)]
    labels_list_train = os.listdir(txt_path_train)
    # labels_list_eval = os.listdir(txt_path_eval)
    lentype=[0 for i in range(class_num)]
    for i in tqdm(labels_list_train):
        file_path = os.path.join(txt_path_train, i)
        file = open(file_path, 'r')  # 打开文件
        file_data = file.readlines()  # 读取所有行
        if len(file_data) == 1:
            lentype[0]+=1
        elif len(file_data) == 2:
            lentype[1]+=1
        else:
            lentype[2]+=1
        for every_row in file_data:
            class_val = every_row.split(' ')[0]
            class_ind = class_list.index(int(class_val))
            class_num_list[class_ind] += 1
        file.close()

    # for j in tqdm(labels_list_eval):
    #     file_path = os.path.join(txt_path_eval, j)
    #     file = open(file_path, 'r')  # 打开文件
    #     file_data = file.readlines()  # 读取所有行
    #     for every_row in file_data:
    #         class_val = every_row.split(' ')[0]
    #         class_ind = class_list.index(int(class_val))
    #         class_num_list[class_ind] += 1
    #     file.close()
    # 输出每一类的数量以及总数
    result=dict(zip(class_categories,class_num_list))
    draw_plot(result,class_num_list)
    # get_type_num(txt_path_train,txt_path_eval,result,class_num_list,folder_path)

    print("-----------------------------------")
    print('total:', sum(class_num_list),lentype)

def get_type_num(txt_path_train,txt_path_eval,result,class_num_list,folder_path):
    type_number=[0,0,0]
    type_num=[1,2,3]
    num_13=0
    for item in os.listdir(txt_path_eval):
        filePath=os.path.join(txt_path_eval,item)
        with open (filePath,'r') as eval_item:  
            first_elements = [int(line.split(' ')[0]) for line in eval_item.readlines()]
            unique_first_elements = list(set(first_elements))
            # if unique_first_elements==[0,2]:
            #     num_13+=1
            #     file_name = os.path.basename(eval_item.name)
            #     image_path = os.path.join(folder_path, 'test', 'images', file_name)
            #     target_image_path = os.path.splitext(image_path)[0] + '.jpg'
            #     shutil.copy(filePath, os.path.join(folder_path,'type_13\labels'))
            #     shutil.copy(target_image_path, os.path.join(folder_path,'type_13\images'))
            # if unique_first_elements==[0]:
            #     file_name = os.path.basename(eval_item.name)
            #     image_path = os.path.join(folder_path, 'test', 'images', file_name)
            #     target_image_path = os.path.splitext(image_path)[0] + '.jpg'
            #     shutil.copy(filePath, os.path.join(folder_path,'type_1\labels'))
            #     shutil.copy(target_image_path, os.path.join(folder_path,'type_1\images'))
            if unique_first_elements==[2]:
                file_name = os.path.basename(eval_item.name)
                image_path = os.path.join(folder_path, 'test', 'images', file_name)
                target_image_path = os.path.splitext(image_path)[0] + '.jpg'
                shutil.copy(filePath, os.path.join(folder_path,'type_3\labels'))
                shutil.copy(target_image_path, os.path.join(folder_path,'type_3\images'))
            type_number[len(unique_first_elements)-1]+=1
    for item in os.listdir(txt_path_train):
        filePath=os.path.join(txt_path_train,item)
        with open (filePath,'r') as eval_item:
            first_elements = [int(line.split(' ')[0]) for line in eval_item.readlines()]
            unique_first_elements = list(set(first_elements))
            # if unique_first_elements==[0,2]:
            #     num_13+=1
            #     file_name = os.path.basename(eval_item.name)
            #     image_path = os.path.join(folder_path, 'train', 'images', file_name)
            #     target_image_path = os.path.splitext(image_path)[0] + '.jpg'
            #     shutil.copy(filePath, os.path.join(folder_path,'type_13\labels'))
            #     shutil.copy(target_image_path, os.path.join(folder_path,'type_13\images'))
            # if unique_first_elements==[0]:
            #     num_13+=1
            #     file_name = os.path.basename(eval_item.name)
            #     image_path = os.path.join(folder_path, 'train', 'images', file_name)
            #     target_image_path = os.path.splitext(image_path)[0] + '.jpg'
            #     shutil.copy(filePath, os.path.join(folder_path,'type_1\labels'))
            #     shutil.copy(target_image_path, os.path.join(folder_path,'type_1\images'))
            if unique_first_elements==[2]:
                num_13+=1
                file_name = os.path.basename(eval_item.name)
                image_path = os.path.join(folder_path, 'train', 'images', file_name)
                target_image_path = os.path.splitext(image_path)[0] + '.jpg'
                shutil.copy(filePath, os.path.join(folder_path,'type_3\labels'))
                shutil.copy(target_image_path, os.path.join(folder_path,'type_3\images'))
            type_number[len(unique_first_elements)-1]+=1
    result_type=dict(zip(type_num,type_number))
    print(num_13)
    draw_plot(result,class_num_list,result_type)

def draw_plot(result:dict,
              class_num_list:list,
            #   result_type:dict,
              ):
    fig,(ax1,ax3)=plt.subplots(2, 1, figsize=(40, 20))
    for name,num in result.items():
        print(name,":",num)
        ax1.bar(name,num,label=name,width=0.5)
    # for type,number in result_type.items():
    #     ax3.bar(type,number,label=type,width=0.5)
    ax1.bar('total',sum(class_num_list),label='total',width=0.5)
    ax1.tick_params(axis='x', labelsize=28)
    ax1.tick_params(axis='y', labelsize=28)
    ax1.set_xlabel('type',fontsize=35)
    ax1.set_ylabel('number',fontsize=35)

    # ax3.tick_params(axis='x', labelsize=28)
    # ax3.tick_params(axis='y', labelsize=28)
    # ax3.set_xlabel('type_number',fontsize=35)
    # ax3.set_ylabel('number',fontsize=35)
    plt.xticks(fontsize=28)
    plt.yticks(fontsize=28)
    plt.legend()
    plt.gca()
    time_now = datetime.now().strftime("%Y%m%d%H%M")
    plt.savefig(f'/ayanjiusheng/project/ultralytics-main/ultralytics/plot_{time_now}.png')
if __name__ == '__main__':
    # 需修改，txt文件所在路径
    txt_path_train = r'D:\CCTSDB2021\all_data\labels'
    txt_path_eval=r'D:\CCTSDB2021\test\labels'
    folder_path=r'D:\CCTSDB2021'
    get_every_class_num(txt_path_train,txt_path_eval,folder_path)